from core.function import Function
from implement.functions.conv.pooling2d_grad import Pooling2DGrad
from utils.functions_conv import im2col_array


class Pooling(Function):
    def __init__(self, kernel_size, stride=1, pad=0):
        """
        池化操作的函数类，用于执行池化操作。

        Args:
            kernel_size (int or tuple of int): 池化核大小。
            stride (int or tuple of int): 步幅。
            pad (int or tuple of int): 填充。
        """
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.pad = pad

    def forward(self, x):
        """
        前向传播，执行池化操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        col = im2col_array(x, self.kernel_size, self.stride, self.pad,
                           to_matrix=False)

        N, C, KH, KW, OH, OW = col.shape
        col = col.reshape(N, C, KH * KW, OH, OW)
        self.indexes = col.argmax(axis=2)
        y = col.max(axis=2)
        return y

    def backward(self, gy):
        """
        反向传播，计算输入变量的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            Variable: 输入变量的梯度。
        """
        return Pooling2DGrad(self)(gy)

    def pooling(self, x):
        """
        执行池化操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        return self(x)
